package vpn

import (
	"context"
	"fmt"
	"github.com/kbinani/screenshot"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
	"log"
	"os/exec"
	"time"
)

type ScreenShare struct {
	displayIndex int
	fps          int
	quality      int
	isStreaming  bool
	CancelFunc   context.CancelFunc
	ctx          context.Context
}

// 创建新的屏幕共享实例
func NewScreenShare() *ScreenShare {
	ctx, cancel := context.WithCancel(context.Background())
	return &ScreenShare{
		displayIndex: 0,  // 默认主屏幕
		fps:          10, // 默认10帧
		quality:      75, // 默认75%质量
		isStreaming:  false,
		CancelFunc:   cancel,
		ctx:          ctx,
	}
}

// 设置显示屏幕索引
func (s *ScreenShare) WithDisplay(index int) *ScreenShare {
	s.displayIndex = index
	return s
}

// 设置帧率
func (s *ScreenShare) WithFPS(fps int) *ScreenShare {
	s.fps = fps
	return s
}

// 设置图像质量
func (s *ScreenShare) WithQuality(quality int) *ScreenShare {
	s.quality = quality
	return s
}

// 修改StartStreaming方法使用H264
func (s *ScreenShare) StartStreaming(peerConnection *webrtc.PeerConnection) error {
	if s.isStreaming {
		return fmt.Errorf("屏幕共享已经在运行")
	}

	// 创建H264视频轨道
	videoTrack, err := webrtc.NewTrackLocalStaticSample(
		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264},
		"screen",
		"stream",
	)
	if err != nil {
		return fmt.Errorf("创建视频轨道失败: %w", err)
	}

	// 添加视频轨道到连接
	_, err = peerConnection.AddTrack(videoTrack)
	if err != nil {
		return fmt.Errorf("添加视频轨道失败: %w", err)
	}

	// 启动屏幕捕获
	s.isStreaming = true

	go s.CaptureAndStreamH264(videoTrack)

	fmt.Printf("🎥 屏幕共享已启动 (屏幕%d, %d FPS)\n", s.displayIndex, s.fps)
	return nil
}

// captureAndStreamH264 使用FFmpeg进行H264编码
func (s *ScreenShare) CaptureAndStreamH264(videoTrack *webrtc.TrackLocalStaticSample) {
	defer func() {
		s.isStreaming = false
		fmt.Println("🛑 屏幕共享已停止")
	}()

	// 使用ffmpeg捕获屏幕（需要安装ffmpeg）
	cmd := exec.Command("ffmpeg",
		//"-f", "x11grab", "-i", ":0.0", // Linux
		// "-f", "avfoundation", "-i", "1", // macOS
		"-f", "gdigrab", "-i", "desktop", // Windows
		"-c:v", "libx264",
		"-preset", "veryfast",
		"-tune", "zerolatency",
		"-f", "h264",
		"-r", "30",
		"-",
	)

	defer cmd.Process.Kill()

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	// 读取H.264数据并发送
	buf := make([]byte, 4096)
	for {
		n, err := stdout.Read(buf)
		if err != nil {
			log.Println("Error reading screen data:", err)
			break
		}
		err = videoTrack.WriteSample(media.Sample{
			Data:     buf[:n],
			Duration: 33 * time.Millisecond, // 30fps
		})
		if err != nil {
			log.Println("Error writing sample:", err)
		}
		log.Println("发送屏幕捕获帧成功")
	}
}

// 停止屏幕共享
func (s *ScreenShare) Stop() {
	if s.CancelFunc != nil {
		s.CancelFunc()
		s.isStreaming = false
	}
}

// 获取屏幕信息
func (s *ScreenShare) GetScreenInfo() string {
	n := screenshot.NumActiveDisplays()
	info := fmt.Sprintf("可用屏幕数量: %d\n", n)

	for i := 0; i < n; i++ {
		bounds := screenshot.GetDisplayBounds(i)
		info += fmt.Sprintf("屏幕%d: %dx%d\n", i, bounds.Dx(), bounds.Dy())
	}

	return info
}
